#include <iostream>
#include <array>

#ifdef WAFER_OPENSSL
#include <openssl/evp.h>
#endif

#include "container.h"

#define MSG_LEN 2
#define MSG_CNT 10
#define INPUT_SIZE (MSG_LEN * MSG_CNT)

#define KEY_LEN 16
#define KEY_CNT 2
#define KEY_SIZE (KEY_LEN * KEY_CNT)

#define AES_DIGEST_LEN MSG_LEN
#define OUTPUT_SIZE 16

enum class AES_TYPE {
  AES128 = 16,
  AES192 = 24,
  AES256 = 32
};

#ifndef WAFER_OPENSSL
extern "C"  {
    void CUDA(MemRef<unsigned char, 1>* inputs, int msgLen, MemRef<unsigned char, 1>* outputs, int cptLen, MemRef<unsigned char, 1>* key, int keyLen, int msgCnt, MemRef<unsigned char, 1>* match_input);
}

void AES_CRACK_OP(MemRef<unsigned char, 1>* inputs, int msgLen, MemRef<unsigned char, 1>* outputs, int cptLen, MemRef<unsigned char, 1>* key, int keyLen, int msgCnt, MemRef<unsigned char, 1>* match_input) {
    CUDA(inputs, msgLen, outputs, cptLen, key, keyLen, msgCnt, match_input);
}
#endif

uint8_t exampleInput[INPUT_SIZE] = {'0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7', '0', '8', '0', '9'};
uint8_t exampleKeys[KEY_SIZE] = {'A', '0', 'B', '1', 'C', '2', 'D', '3', 
                                'E', '4', 'F', '5', 'G', '6', 'H', '7', 
                                'I', '8', 'J', '9', 'K', '0', 'L', '1', 
                                'M', '2', 'N', '3', 'O', '4', 'P', '5',};

MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> keys{exampleKeys, KEY_SIZE};
// output: 99
uint8_t exampleOutput[OUTPUT_SIZE] = {
0x3c, 0x2c, 0x9b, 0xc6, 0xc2, 0xf5, 0x62, 0x7d, 0xb3, 0x0f, 0x93, 0x40, 0xef, 0xa9, 0xc7, 0x2a
};
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};

uint8_t _matched_input[MSG_LEN] = {0, 0};
MemRef<unsigned char, 1> matched_input{_matched_input, MSG_LEN};


void printMatches(const char* title, uint8_t* match) {
  // Print the output.
  std::cout << std::endl;
  std::cout << title;
  std::cout << "[ ";
  for (size_t i = 0; i < MSG_LEN; i++) {
    printf("%c ", match[i]);
  }
  std::cout << "]";
  std::cout << std::endl;
}


void generateResultGPUAES() {
  int keyLen = static_cast<int>(AES_TYPE::AES128);
  
#ifdef WAFER_OPENSSL
  std::cout << "--------------------------------------------------------\nOpenSSL AES\n";

  const clock_t begin_time = clock();

for (int i = 0; i < MSG_CNT; i++) {
  printInput(i, exampleInput + (i * MSG_LEN));

  EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();

  const EVP_CIPHER *cipher;
  if (keyLen == 16)
    cipher = EVP_aes_128_ecb();
  else if (keyLen == 24)
    cipher = EVP_aes_192_ecb();
  else // keyLen == 32
    cipher = EVP_aes_256_ecb();

  EVP_DecryptInit_ex(ctx, cipher, NULL, exampleKey, NULL);
  int outlen;
  int plaintext_len = 0;
  EVP_DecryptUpdate(ctx, exampleOutput + (i * AES_DIGEST_LEN), &outlen, exampleInput + (i * MSG_LEN), AES_DIGEST_LEN);
  plaintext_len += outlen;
  EVP_DecryptFinal_ex(ctx, exampleOutput + (i * AES_DIGEST_LEN) + plaintext_len, &outlen);
  plaintext_len += outlen;

  printResult(exampleOutput + (i * AES_DIGEST_LEN));
  EVP_CIPHER_CTX_free(ctx);
}

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#else
  const clock_t begin_time = clock();
  #pragma wafer parallel
  AES_CRACK_OP(&input, MSG_LEN, &output, OUTPUT_SIZE, &keys, keyLen, MSG_CNT, &matched_input);
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  printMatches("MLIR AES crack result: ", _matched_input);
  
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#endif
}

int main() {
  generateResultGPUAES();
  return 0;
}
